#!/usr/bin/env python3
"""
Script: run_ionization.py
Function: Run TOPAS simulations for ionization
          with carbon ion beam energy from 1 MeV to 5400 MeV
Author: Converted from bash script
"""

import os
import shutil
import subprocess
import sys
from pathlib import Path


def main():
    # Colors for output (Windows compatible)
    RED = '\033[0;31m'
    GREEN = '\033[0;32m'
    YELLOW = '\033[1;33m'
    BLUE = '\033[0;34m'
    NC = '\033[0m'  # No Color

    # Check if carbon1.txt exists
    if not os.path.exists("carbon1.txt"):
        print(f"{RED}Error: carbon1.txt file not found!{NC}")
        sys.exit(1)

    # Check if TOPAS executable exists
    topas_path = "../bin/topas"
    if not os.path.exists(topas_path):
        print(f"{RED}Error: {topas_path} executable not found!{NC}")
        sys.exit(1)

    # Create results1 directory
    os.makedirs("results1", exist_ok=True)

    print("Starting batch TOPAS simulations...")
    print("Energy range: 1 MeV to 5400 MeV")
    print("=" * 50)

    # Run simulations for different energies
    for energy in range(1, 5401):
        print(f"Running simulation for energy {energy} MeV...")

        # Create temporary config file
        temp_file = f"carbon1_{energy}MeV.txt"

        # Read original config file
        try:
            with open("carbon1.txt", 'r', encoding='utf-8') as f:
                content = f.read()
        except Exception as e:
            print(f"{RED}Error reading carbon1.txt: {e}{NC}")
            continue

        # Replace energy and output filename
        # Original: d:So/Demo/BeamEnergy = 1800 MeV
        content = content.replace("d:So/Demo/BeamEnergy = 1800 MeV",
                                 f"d:So/Demo/BeamEnergy = {energy} MeV")

        # Original: s:Sc/Dose1/OutputFile = "ionization"
        content = content.replace('s:Sc/Dose1/OutputFile = "ionization"',
                                 f's:Sc/Dose1/OutputFile = "results1/ionization_{energy}"')

        # Write temporary config file
        try:
            with open(temp_file, 'w', encoding='utf-8') as f:
                f.write(content)
        except Exception as e:
            print(f"{RED}Error writing temporary file {temp_file}: {e}{NC}")
            continue

        # Run TOPAS
        try:
            result = subprocess.run([topas_path, temp_file],
                                  capture_output=True, text=True, cwd=".")
            if result.returncode == 0:
                print(f"{GREEN}✓ Energy {energy} MeV simulation completed{NC}")
            else:
                print(f"{RED}✗ Energy {energy} MeV simulation failed{NC}")
                if result.stderr:
                    print(f"Error output: {result.stderr}")
        except Exception as e:
            print(f"{RED}✗ Energy {energy} MeV simulation failed: {e}{NC}")

        # Clean up temporary file
        try:
            os.remove(temp_file)
        except OSError:
            pass  # File might not exist, that's okay

        # Show progress every 100 energies
        if energy % 100 == 0:
            progress_percent = round(energy * 100 / 5400, 1)
            print(f"Progress: {energy}/5400 ({progress_percent}%)")

    print("=" * 50)
    print("All simulations completed!")
    print("Result files saved in results1/ directory")
    print("Total of 5400 simulations run for different energies")


if __name__ == "__main__":
    main()
